﻿using System;
using System.Collections.Generic;
using System.Text.RegularExpressions;
using System.IO;
using System.Diagnostics;
using System.Globalization;

namespace StreamSplit
{
    class Program
    {
        static void Main(string[] args)
        {
            Console.WriteLine("-StreamSplit 0.2-");
            Console.WriteLine();

            // check mkvtoolnix
            if (!File.Exists(AppDomain.CurrentDomain.BaseDirectory + "mkvmerge.exe") || !File.Exists(AppDomain.CurrentDomain.BaseDirectory + "mkvextract.exe"))
            {
                Console.WriteLine("mkvmerge.exe and/or mkvextract.exe not found.");
                Console.ReadLine();
                Environment.Exit(0);
            }
            // end check mkvtoolnix

            // check arguments
            if (args.Length < 2) // insufficent arguments or input does not exist
            {
                printhelp("Insufficient arguments.");
            }
            if (!File.Exists(args[0]) || !File.Exists(args[1]))
            {
                printhelp("Input doesn't exist.");
            }
            string stream_path = null;
            string avs_path = null;
            decimal fps = 0;
            int delay = 0;
            bool delay_set = false;
            string sbr = null;

            if (Path.GetExtension(args[0]) == ".avs")
            {
                avs_path = args[0];
                stream_path = args[1];
            }
            else if (Path.GetExtension(args[1]) == ".avs")
            {
                avs_path = args[1];
                stream_path = args[0];
            }
            else
            {
                printhelp("No .avs script specified.");
            }
            for (int i = 2; i < args.Length; i++)
            {
                if (args[i] == "--fps")
                {
                    if (args.Length > i)
                    {
                        fps = process_fps(args[i + 1]);
                        if (fps == 0)
                        {
                            printhelp("Invalid fps specified.");
                        }
                        i++;
                    }
                    else
                    {
                        printhelp("No value for --fps specified");
                    }
                }
                else if (args[i] == "--delay")
                {
                    if (args.Length > i)
                    {
                        if(!Int32.TryParse(args[i + 1], out delay))
                        {
                            printhelp("Invalid delay specified.");
                        }
                        i++;
                        delay_set = true;
                    }
                    else
                    {
                        printhelp("No value for --delay specified.");
                    }
                }
                else if (args[i] == "--sbr")
                {
                    sbr = " --aac-is-sbr 0";
                }
            }

            // post argument fps handling
            string input_fps;
            if (fps == 0)
            {
                Console.WriteLine("Arguments didn't contain fps, please enter them now:");
                input_fps = Console.ReadLine();
                fps = process_fps(input_fps);
                if (fps == 0)
                {
                    printhelp("Invalid fps specified.");
                }
            }
            // fps handling end

            Console.WriteLine("Using fps: " + fps);

            // get delay from audio file in format "DELAY Xms"
            if (!delay_set)
            {
                Console.WriteLine("No delay specified. Checking stream file name...");
                Regex delay_regex = new Regex(@"DELAY -?\d+ms");
                Match delay_match = delay_regex.Match(stream_path);
                if (delay_match.Success)
                {
                    delay = Convert.ToInt32(delay_match.Value.Substring(6, delay_match.Value.Length - 8));
                    Console.WriteLine("Found delay: " + delay + ".");
                }
                else
                {
                    Console.WriteLine("No delay found.");
                }
            }
            // get delay end

            Console.WriteLine("Using delay: " + delay);

            // get Trim parameter from avs file
            List<long> start_frames = new List<long>();
            List<long> end_frames = new List<long>();
            StreamReader avs = new StreamReader(avs_path);
            bool found_trim = false;
            string avs_line;
            while ((avs_line = avs.ReadLine()) != null && !found_trim)
            {
                string clean_line = avs_line.Replace(" ", null).Replace("++", "+"); // remove spaces and replace ++ with + in avs line
                Regex trim_line_regex = new Regex(@"Trim\(\d+\,\d+\)(\+Trim\(\d+\,\d+\))*", RegexOptions.IgnoreCase);
                Match trim_line_match = trim_line_regex.Match(clean_line);
                if (trim_line_match.Success)
                {
                    string[] trims = trim_line_match.Value.Split('+');
                    for (int i = 0; i < trims.Length; i++)
                    {
                        Regex frame_regex = new Regex(@"\d+");
                        Match frame_match = frame_regex.Match(trims[i]);
                        start_frames.Add(Convert.ToInt64(frame_match.Value));
                        frame_match = frame_match.NextMatch();
                        end_frames.Add(Convert.ToInt64(frame_match.Value));
                    }
                    found_trim = true;
                    Console.WriteLine("Using Trim found in this line:");
                    Console.WriteLine(avs_line);
                }
            }
            avs.Close();
            if (!found_trim)
            {
                Console.WriteLine("No Trim found.");
                Console.ReadLine();
                Environment.Exit(0);
            }
            // get trim end

            // create mkvmerge split line
            string split_parameter = String.Empty;
            for (int i = 0; i < start_frames.Count; i++)
            {
                long start_ms = (long)Math.Round(1000 * start_frames[i] / fps, 0);
                long end_ms = (long)Math.Round(1000 * end_frames[i] / fps, 0);
                TimeSpan start_span = TimeSpan.FromMilliseconds(start_ms);
                TimeSpan end_span = TimeSpan.FromMilliseconds(end_ms);
                string start_timestamp = string.Format("{0:D2}:{1:D2}:{2:D2}.{3:D3}", start_span.Hours, start_span.Minutes, start_span.Seconds, start_span.Milliseconds);
                string end_timestamp = string.Format("{0:D2}:{1:D2}:{2:D2}.{3:D3}", end_span.Hours, end_span.Minutes, end_span.Seconds, end_span.Milliseconds);
                split_parameter = split_parameter + start_timestamp + "-" + end_timestamp + ",+";
            }
            split_parameter = split_parameter.Substring(0, split_parameter.Length - 2); // delete the last comma and plus
            Console.WriteLine();
            Console.WriteLine("Starting mkvmerge:");

            // execute mkvmerge
            Process mkvmerge = new Process();
            mkvmerge.StartInfo.UseShellExecute = false; // execute in current window
            mkvmerge.StartInfo.FileName = "mkvmerge.exe";
            mkvmerge.StartInfo.Arguments = "--ui-language en -o \"" + Path.GetFileNameWithoutExtension(avs_path) + ".mkv\"" + sbr + " --split parts:" + split_parameter + " -y 0:" + delay + " \"" + stream_path + "\"";
            mkvmerge.Start();
            mkvmerge.WaitForExit();
            Console.WriteLine("Complete.");
            Console.WriteLine();
            Console.WriteLine("Starting mkvextract:");
            // execute mkvmerge end

            // execute mkvextract
            Process mkvextract = new Process();
            mkvextract.StartInfo.UseShellExecute = false; // execute in current window
            mkvextract.StartInfo.FileName = "mkvextract.exe";
            mkvextract.StartInfo.Arguments = "--ui-language en tracks \"" + Path.GetFileNameWithoutExtension(avs_path) + ".mkv\" 0:\"" + Path.GetFileNameWithoutExtension(avs_path) + Path.GetExtension(stream_path) + "\"";
            mkvextract.Start();
            mkvextract.WaitForExit();
            File.Delete(Path.GetFileNameWithoutExtension(avs_path) + ".mkv"); // delete temp mkv
            Console.WriteLine("Complete.");
            Console.WriteLine();
            Console.WriteLine("Well, let's hope it worked. Press any key to exit...");
            Console.ReadLine();
            // execute mkvextract end
        }

        static void printhelp(string error)
        {
            Console.WriteLine(error);
            Console.WriteLine();
            Console.WriteLine("Example usage:");
            Console.WriteLine("   StreamSplit.exe script.avs audio.m4a --fps 30000/1001");
            Console.WriteLine("You can also drag-and-drop the Avisynth script and the input Stream onto StreamSplit.exe");
            Console.WriteLine("Options:");
            Console.WriteLine("   --fps <float|rational>  Specify frame rate");
            Console.WriteLine("   --delay <integer>       Specify stream delay in ms (overrides value found in file name)");
            Console.WriteLine("   --sbr                   Tell mkvmerge to treat stream as SBR AAC");
            Console.ReadLine();
            Environment.Exit(0);
        }

        static decimal process_fps(string input)
        {
            Regex fps_float_regex = new Regex(@"^\d+(\.\d+)?$");
            Match fps_float_match = fps_float_regex.Match(input);
            Regex fps_rational_regex = new Regex(@"^\d+/\d+$");
            Match fps_rational_match = fps_rational_regex.Match(input);
            if (fps_float_match.Success)
            {
                return decimal.Parse(input, CultureInfo.GetCultureInfo("en-US"));
            }
            else if (fps_rational_match.Success)
            {
                string[] rational = input.Split('/');
                return decimal.Parse(rational[0]) / Int32.Parse(rational[1]);
            }
            else
            {
                return 0;
            }
        }
    }
}